package com.my.netty.nio.net.selector;

import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.charset.Charset;
import java.util.Iterator;
import java.util.Set;

@Slf4j
public class SelectorNioWriteServer {

    public static void main(String[] args) throws IOException {
        //NIO的非阻塞式TCP服务端+Selector（单线程处理）+ 演示可写事件

        ServerSocketChannel ssc = ServerSocketChannel.open();
        ssc.bind(new InetSocketAddress(8080));
        ssc.configureBlocking(false);

        Selector selector = Selector.open();
        ssc.register(selector, SelectionKey.OP_ACCEPT, null);

        while (true) {
            log.debug("Selector等待事件...");
            selector.select();

            Set<SelectionKey> selectionKeys = selector.selectedKeys();
            Iterator<SelectionKey> iter = selectionKeys.iterator();
            while (iter.hasNext()) {
                SelectionKey key = iter.next();

                if (key.isAcceptable()) {
                    SocketChannel sc = ssc.accept();
                    sc.configureBlocking(false);
                    SelectionKey scKey = sc.register(selector, 0, null);
                    log.debug("获取客户端连接:{}", sc.getRemoteAddress());

                    //模拟向客户端发送大量数据
                    ByteBuffer buffer = getBigByteBuffer();

                    /*
                    //不使用selector 不一定能一次写完，需要循环
                    while (buffer.hasRemaining()) {
                        //返回一次写的字节数 (如果发送的数据量非常大，write一次无法写入所有数据，网络缓冲区可能被写满，此时再write也不能写入数据，返回值为0)
                        int len = sc.write(buffer);
                        log.debug("向客户端 {} 写入 {} 字节", sc.getRemoteAddress(), len);
                    }
                     */

                    //使用selector，先写一次，如果没写完则把channel注册到selector，等待下次可写事件时，再写剩余的数据
                    int len = sc.write(buffer);
                    log.debug("向客户端 {} 写入 {} 字节", sc.getRemoteAddress(), len);
                    if (buffer.hasRemaining()) {
                        //添加关注可写事件，注意通过或运算，避免覆盖原有关注事件
                        scKey.interestOps(scKey.interestOps() | SelectionKey.OP_WRITE);
                        //把未写完的buffer放到附件中，下次可写事件时使用
                        scKey.attach(buffer);
                    }

                } else if (key.isWritable()) {
                    //可写事件
                    SocketChannel sc = (SocketChannel) key.channel();
                    //获取之前未写完的buffer
                    ByteBuffer buffer = (ByteBuffer) key.attachment();
                    //这里没写完也不要紧，在while循环下次可写事件中还会继续写剩余的数据
                    int len = sc.write(buffer);
                    log.debug("向客户端 {} 写入 {} 字节", sc.getRemoteAddress(), len);
                    //如果已经写完，需要释放buffer的内存空间；并取消关注可写事件，否则出现NPE
                    if (!buffer.hasRemaining()) {
                        key.attach(null);
                        key.interestOps(key.interestOps() & (~SelectionKey.OP_WRITE));
                    }
                }

                iter.remove();
            }
        }

    }

    private static ByteBuffer getBigByteBuffer() {
        //模拟向客户端发送大量数据
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < 30_000_000; i++) {
            char c = (char) ('a' + ('z' - 'a' + 1) * Math.random());
            sb.append(c);
        }
        return Charset.defaultCharset().encode(sb.toString());
    }

}
